﻿//
// $HeadURL: http://codeka-memcache.googlecode.com/svn/trunk/Codeka.Memcache/ServerConnection.cs $
// $Revision: 7 $
// $Author: dean.codeka $     $Date: 2010-02-20 12:24:59 +0000 (Sat, 20 Feb 2010) $
//
// Copyright © 2009-2010 Dean Harding, all rights reserved.
//

using Codeka.Memcache.Config;
using System.Net.Sockets;
using System;
using Codeka.Memcache.Commands;
using System.Collections.Generic;
using System.IO;
using System.Threading;

namespace Codeka.Memcache
{
    /// <summary>
    /// This class represents a connection to a server. These are stored in
    /// the <see cref="ServerPool"/>, once instance per server.
    /// </summary>
    public class ServerConnection
    {
        private ServerConfig config;
        private Socket s;
        private bool inProgress;
        private Queue<ExecuteContext> queuedOperations;
        private byte[] receiveBuffer;
        private MemoryStream response;

        /// <summary>
        /// Constructs a new <see cref="ServerConnection"/> that'll use the given <see cref="ServerConfig"/>
        /// to connect to the server.
        /// </summary>
        public ServerConnection(ServerConfig config)
        {
            this.config = config;
            this.queuedOperations = new Queue<ExecuteContext>();
            this.receiveBuffer = new byte[4096];
            this.response = new MemoryStream();
        }

        /// <summary>
        /// Begins executing the given <see cref="Command"/>. Because we execute commands
        /// one at a time, this may actually queue up the command to be executed, and callers
        /// will be able to query the <see cref="Command"/>'s <b>IsQueued</b> property to
        /// know.
        /// </summary>
        /// <param name="cmd">The <see cref="Command"/> to execute.</param>
        /// <param name="callback">An <see cref="AsyncCallback"/> to call when the command has
        /// finished executing.</param>
        /// <param name="state">A <c>state</c> object to pass back to the <see cref="AsyncCallback"/>
        /// when the command finishes.</param>
        /// <returns>An <see cref="IAsyncResult"/> instance you can use to determine the state
        /// of the operation.</returns>
        public IAsyncResult BeginExecuteCommand(Command cmd, AsyncCallback callback, object state)
        {
            lock (this)
            {
                EnsureOpen();

                var context = new ExecuteContext {
                        AsyncState = state,
                        Command = cmd,
                        Callback = callback,
                    };
                
                cmd.IsQueued = true;
                queuedOperations.Enqueue(context);

                // if there's nothing currently in progress, just start this one now
                if (!inProgress)
                {
                    ExecuteNextCommand();
                }

                return context;
            }
        }

        /// <summary>
        /// Waits for the execution of a command to complete.
        /// </summary>
        public void EndExecuteCommand(IAsyncResult res)
        {
            var context = (ExecuteContext)res;
            if (!context.IsCompleted)
                context.AsyncWaitHandle.WaitOne();

            if (context.Exception != null)
                throw new Exception("An exception occured during an ExecuteCommand", context.Exception);
        }

        /// <summary>
        /// This is called when a command is placed on the queue and we're not
        /// doing anything, or when a command finishes and we want to execute the
        /// next one.
        /// </summary>
        private void ExecuteNextCommand()
        {
            // note: we assume we're already in a lock(this)

            if (queuedOperations.Count == 0)
            {
                inProgress = false;
                return;
            }

            // get the next command and mark ourselves inProgress
            var context = queuedOperations.Dequeue();
            inProgress = true;

            // serialize all of the packets to a memory stream
            var ms = new MemoryStream();
            var writer = new BinaryWriter(ms);
            foreach (var pkt in context.Command.GetPackets())
            {
                pkt.Serialize(writer);
            }
            writer.Flush();

            // start sending the data!
            s.BeginSend(ms.GetBuffer(), 0, (int) ms.Length, SocketFlags.None, OnSendComplete, context);
        }

        /// <summary>
        /// This is called when a send operation has completed. We have to wait for the response.
        /// </summary>
        void OnSendComplete(IAsyncResult res)
        {
            var context = res.AsyncState as ExecuteContext;
            if (context == null)
                return;

            try
            {
                int numbytes = s.EndSend(res);
                context.CompletedSynchronously = context.CompletedSynchronously && res.CompletedSynchronously;

                // log.Write(sent)

                // simply wait for the response (reinitialise the current response as well)...
                response.SetLength(0);

                s.BeginReceive(receiveBuffer, 0, receiveBuffer.Length, SocketFlags.None, OnReceiveComplete, context);
            }
            catch (Exception e)
            {
                context.Exception = e;
                context.MarkComplete(res.CompletedSynchronously);
            }
        }

        /// <summary>
        /// This is called when a receive operation has finished. We've got some data so add
        /// it to the current buffer, then check if we've got a complete response and process
        /// it if so.
        /// </summary>
        void OnReceiveComplete(IAsyncResult res)
        {
            var context = res.AsyncState as ExecuteContext;
            if (context == null)
                return;

            try
            {
                // append what we've just received to the "response" buffer
                int numbytes = s.EndReceive(res);
                context.CompletedSynchronously = context.CompletedSynchronously && res.CompletedSynchronously;

                response.Write(receiveBuffer, 0, numbytes);

                // if we've got at least a PacketHeader, then begin by deserializing it...
                if (response.Length >= PacketHeader.SerializedLength)
                {
                    // create a whole new MemoryStream to process it (so that we don't
                    // mess with the one we've got)
                    var ms = new MemoryStream(response.GetBuffer());
                    var reader = new BinaryReader(ms);

                    // deserialize the header so we can find out how big the rest of it is
                    // supposed to be
                    var header = new PacketHeader();
                    header.Deserialize(reader);

                    // if we've got enough data, then deserialize the rest of the packet
                    if (response.Length - PacketHeader.SerializedLength >= header.TotalBodyLength)
                    {
                        var packet = new Packet();
                        packet.Deserialize(header, reader);

                        // we've got a complete packet so save off any packet data that might
                        // be left over, then call back the command to tell it we've got a 
                        // packet for it...
                        var remaining = reader.ReadBytes((int)(ms.Length - ms.Position));
                        response = new MemoryStream();
                        if (remaining.Length > 0)
                            response.Write(remaining, 0, remaining.Length);

                        if (!context.Command.ProcessResponse(packet))
                        {
                            context.MarkComplete(res.CompletedSynchronously);

                            // if this was the last packet then we're finished processing!
                            lock (this)
                            {
                                inProgress = false;
                                ExecuteNextCommand();
                                return;
                            }
                        }
                    }
                }

                // if we get here, then we've got to start another read operation -
                // there's still more work to be done!
                s.BeginReceive(receiveBuffer, 0, receiveBuffer.Length, SocketFlags.None, OnReceiveComplete, context);
            }
            catch (Exception e)
            {
                context.Exception = e;
                context.MarkComplete(res.CompletedSynchronously);
            }
        }

        /// <summary>
        /// This is called to make sure the connection is open and ready to
        /// process command(s).
        /// </summary>
        private void EnsureOpen()
        {
            if (s != null)
                return;

            s = new Socket(config.EndPoint.AddressFamily, SocketType.Stream, ProtocolType.Tcp);
            s.Connect(config.EndPoint);
        }

        /// <summary>
        /// This is the <see cref="IAsyncResult"/> we'll return from the operation and represents the
        /// state we'll need to finish the operation.
        /// </summary>
        class ExecuteContext : IAsyncResult
        {
            public ExecuteContext()
            {
                AsyncWaitHandle = new EventWaitHandle(false, EventResetMode.ManualReset);
                CompletedSynchronously = true; // this is the default...
            }

            /// <summary>
            /// Gets or sets the callback we'll call when the command has finished executing.
            /// </summary>
            public AsyncCallback Callback { get; set; }

            /// <summary>
            /// The async state object that the user passes in when they call <b>BeginExecute</b>.
            /// </summary>
            public object AsyncState { get; set; }

            /// <summary>
            /// A <see cref="WaitHandle"/> that you can use to wait for the operation to complete.
            /// </summary>
            public WaitHandle AsyncWaitHandle { get; private set; }

            /// <summary>
            /// A value which indicates whether the operation completed on the calling thread.
            /// </summary>
            public bool CompletedSynchronously { get; set; }

            /// <summary>
            /// Gets a value which indicates whether we're finished.
            /// </summary>
            public bool IsCompleted { get; set; }

            /// <summary>
            /// Gets or sets the <see cref="Command"/> we're executing.
            /// </summary>
            public Command Command { get; set; }

            /// <summary>
            /// If an exception happens in one our callbacks, we'll want to 'cancel'
            /// the operation and throw an exception from <see cref="EndExecuteCommand"/>.
            /// </summary>
            public Exception Exception { get; set; }

            /// <summary>
            /// When the operation is complete, we want to set <b>IsCompleted</b>, set the
            /// event and so on.
            /// </summary>
            public void MarkComplete(bool completedSynchronously)
            {
                CompletedSynchronously = CompletedSynchronously && completedSynchronously;
                IsCompleted = true;
                ((EventWaitHandle)AsyncWaitHandle).Set();

                // if there's a callback, call it now
                if (Callback != null)
                    Callback(this);
            }
        }
    }
}
